"""知识图谱构建和问答核心模块

该模块提供了从文件中提取知识三元组、构建知识图谱以及基于知识图谱进行问答的核心功能。
支持从多种文件格式（txt、md、csv、json、pdf、doc、docx）中提取知识，并使用LlamaIndex构建知识图谱索引。

主要功能：
1. 从文本中提取知识三元组
2. 从文件构建知识图谱
3. 基于知识图谱进行问答

依赖：
- llama_index：用于构建知识图谱和问答引擎
- 自定义LLM配置：用于问答生成
- PyPDF2：用于处理PDF文件
- python-docx：用于处理DOCX文件
- docx2pdf：用于处理DOC文件转换
- tika-python：用于处理多种文档格式
- python-magic：用于文件类型检测
"""

from pathlib import Path
from typing import List, Dict, Union
import json, csv, re, io, os
from llama_index.core import KnowledgeGraphIndex, StorageContext, Document
from llama_index.core.graph_stores import SimpleGraphStore
from llama_index.core.node_parser import SimpleNodeParser
from config.llm import deepseek_llm  # 自定义LLM配置
from utils.logger import get_logger  # 导入日志模块

# 初始化日志记录器
logger = get_logger('rag_core')

# 导入处理 PDF 和 DOC 文件的库
try:
    import PyPDF2  # 处理 PDF 文件
    import docx  # 处理 DOCX 文件
    from docx2pdf import convert  # 处理 DOC 文件转换
    import tika  # 处理多种文档格式
    from tika import parser as tika_parser
    import magic  # 文件类型检测
    
    # 初始化 tika
    tika.initVM()
    
    DOC_PDF_SUPPORT = True
except ImportError:
    DOC_PDF_SUPPORT = False
    logger.warning("未安装 PyPDF2、python-docx、docx2pdf、tika-python 或 python-magic 库，无法处理 PDF 和 DOC 文件。")
    logger.info("请运行: pip install PyPDF2 python-docx docx2pdf tika-python python-magic")

LLM = deepseek_llm()  # 初始化LLM模型实例

TRIPLE_PATTERN = re.compile(r"(.+?)[\\s,，]*(是|属于|提出|拥有|包括)[\\s,，]*(.+)")

def detect_file_type(file_path: Path) -> str:
    """
    使用 python-magic 检测文件的真实类型。
    
    参数:
        file_path (Path): 文件路径
        
    返回:
        str: 文件的MIME类型
    """
    if not DOC_PDF_SUPPORT or not magic:
        return "unknown/unknown"
        
    try:
        mime = magic.Magic(mime=True)
        file_type = mime.from_file(str(file_path))
        logger.debug(f"文件 {file_path.name} 的MIME类型: {file_type}")
        return file_type
    except Exception as e:
        logger.warning(f"使用 python-magic 检测文件类型失败: {str(e)}，将根据扩展名猜测")
        # 如果检测失败，返回基于扩展名的猜测
        ext_map = {
            ".txt": "text/plain",
            ".md": "text/markdown",
            ".csv": "text/csv",
            ".json": "application/json",
            ".pdf": "application/pdf",
            ".doc": "application/msword",
            ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
        }
        file_type = ext_map.get(file_path.suffix.lower(), "application/octet-stream")
        logger.debug(f"根据扩展名猜测文件 {file_path.name} 的MIME类型: {file_type}")
        return file_type

def _extract_triples_from_line(line: str) -> List[Dict[str, str]]:
    """
    从文本行中提取知识三元组（主体-关系-客体）。
    
    参数:
        line (str): 需要解析的文本行
        
    返回:
        List[Dict[str, str]]: 包含提取出的三元组的列表，每个三元组是一个字典，
                              包含'source'(主体)、'relation'(关系)和'target'(客体)键。
                              如果没有找到三元组，则返回空列表。
    """
    m = TRIPLE_PATTERN.search(line)
    if m:
        subj, rel, obj = m.groups()
        return [{"source": subj.strip(), "relation": rel.strip(), "target": obj.strip()}]
    return []

def _extract_text_from_pdf(file_path: Path) -> List[str]:
    """
    从PDF文件中提取文本内容。
    
    参数:
        file_path (Path): PDF文件路径
        
    返回:
        List[str]: 提取出的文本行列表
    """
    if not DOC_PDF_SUPPORT:
        return [f"无法处理PDF文件 {file_path.name}，请安装必要的库"]
    
    try:
        text_lines = []
        with open(file_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                text = page.extract_text()
                if text:
                    text_lines.extend(text.splitlines())
        return text_lines
    except Exception as e:
        logger.error(f"提取PDF文本时出错: {str(e)}")
        return [f"处理PDF文件时出错: {str(e)}"]

def _extract_text_from_docx(file_path: Path) -> List[str]:
    """
    从DOCX文件中提取文本内容。
    
    参数:
        file_path (Path): DOCX文件路径
        
    返回:
        List[str]: 提取出的文本行列表
    """
    if not DOC_PDF_SUPPORT:
        return [f"无法处理DOCX文件 {file_path.name}，请安装必要的库"]
    
    try:
        doc = docx.Document(file_path)
        return [para.text for para in doc.paragraphs if para.text.strip()]
    except Exception as e:
        return [f"处理DOCX文件时出错: {str(e)}"]

def _extract_text_from_doc(file_path: Path) -> List[str]:
    """
    从DOC文件中提取文本内容。
    
    参数:
        file_path (Path): DOC文件路径
        
    返回:
        List[str]: 提取出的文本行列表
    """
    if not DOC_PDF_SUPPORT:
        return [f"无法处理DOC文件 {file_path.name}，请安装必要的库"]
    
    try:
        # 使用tika解析文档
        parsed = tika_parser.from_file(str(file_path))
        if parsed and 'content' in parsed and parsed['content']:
            text = parsed['content']
            return text.splitlines()
        else:
            return [f"无法从DOC文件 {file_path.name} 提取内容"]
    except Exception as e:
        logger.error(f"提取DOC文本时出错: {str(e)}")
        return [f"处理DOC文件时出错: {str(e)}"]

def build_kg_from_file(file_path: Union[str, Path]) -> List[Dict[str, str]]:
    """
    从文件中构建知识图谱三元组。
    
    根据文件类型（txt、md、csv、json、pdf、doc、docx）读取内容并提取知识三元组。
    使用文件MIME类型和扩展名双重检测，提高处理准确性。
    如果没有找到任何三元组，会创建一个默认三元组。
    
    参数:
        file_path (Union[str, Path]): 文件路径，可以是字符串或Path对象
        
    返回:
        List[Dict[str, str]]: 包含所有提取出的知识三元组的列表
    """
    file_path = Path(file_path)
    triples = []
    
    # 检测文件类型
    mime_type = detect_file_type(file_path)
    file_ext = file_path.suffix.lower()
    
    # 处理文本文件
    if mime_type.startswith("text/") or file_ext in {".txt", ".md"}:
        try:
            lines = file_path.read_text(encoding="utf-8", errors="ignore").splitlines()
            for line in lines:
                triples.extend(_extract_triples_from_line(line.strip()))
        except Exception as e:
            logger.error(f"处理文本文件时出错: {str(e)}")
    
    # 处理CSV文件
    elif mime_type == "text/csv" or file_ext == ".csv":
        try:
            with file_path.open("r", encoding="utf-8", errors="ignore") as f:
                for row in csv.reader(f):
                    for cell in row:
                        triples.extend(_extract_triples_from_line(cell))
        except Exception as e:
            logger.error(f"处理CSV文件时出错: {str(e)}")
    
    # 处理JSON文件
    elif mime_type == "application/json" or file_ext == ".json":
        try:
            data = json.loads(file_path.read_text(encoding="utf-8", errors="ignore"))
            if isinstance(data, list):
                for item in data:
                    text = item if isinstance(item, str) else item.get("text", "")
                    triples.extend(_extract_triples_from_line(text))
        except Exception as e:
            logger.error(f"处理JSON文件时出错: {str(e)}")
    
    # 处理PDF文件
    elif mime_type == "application/pdf" or file_ext == ".pdf":
        lines = _extract_text_from_pdf(file_path)
        for line in lines:
            triples.extend(_extract_triples_from_line(line.strip()))
    
    # 处理DOCX文件
    elif mime_type.endswith("wordprocessingml.document") or file_ext == ".docx":
        lines = _extract_text_from_docx(file_path)
        for line in lines:
            triples.extend(_extract_triples_from_line(line.strip()))
    
    # 处理DOC文件
    elif mime_type == "application/msword" or file_ext == ".doc":
        lines = _extract_text_from_doc(file_path)
        for line in lines:
            triples.extend(_extract_triples_from_line(line.strip()))
    
    # 尝试使用Tika处理其他未知类型文件
    elif DOC_PDF_SUPPORT:
        try:
            parsed = tika_parser.from_file(str(file_path))
            if parsed and 'content' in parsed and parsed['content']:
                text = parsed['content']
                for line in text.splitlines():
                    triples.extend(_extract_triples_from_line(line.strip()))
        except Exception as e:
            logger.error(f"使用Tika处理文件时出错: {str(e)}")
    
    # 如果没有提取到三元组，创建一个默认三元组
    if not triples:
        logger.warning(f"未从文件 {file_path.name} 中提取到三元组，使用默认三元组")
        triples.append({"source": file_path.stem, "relation": "包含", "target": "内容"})
    else:
        logger.info(f"从文件 {file_path.name} 中提取到 {len(triples)} 个三元组")
        
    return triples

def _build_lindex(triples: List[Dict[str, str]]):
    """
    根据三元组构建LlamaIndex知识图谱索引。
    
    将三元组转换为句子，然后创建Document对象，最后构建知识图谱索引。
    
    参数:
        triples (List[Dict[str, str]]): 知识三元组列表
        
    返回:
        KnowledgeGraphIndex: 构建好的知识图谱索引对象
    """
    sentences = [f"{t['source']} {t['relation']} {t['target']}" for t in triples]
    # 使用Document对象而不是字典
    docs = [Document(text=s) for s in sentences]
    nodes = SimpleNodeParser().get_nodes_from_documents(docs)
    graph_store = SimpleGraphStore()
    storage = StorageContext.from_defaults(graph_store=graph_store)
    return KnowledgeGraphIndex(nodes=nodes, storage_context=storage, llm=LLM)

def ask_question_with_kg(question: str, triples: List[Dict[str, str]] = None) -> str:
    """使用知识图谱回答问题
    
    参数:
        question (str): 问题
        triples (List[Dict[str, str]], optional): 三元组列表. Defaults to None.
        
    返回:
        str: 回答
    """
    logger.info(f"收到问题: {question}")
    
    if not triples:
        # 如果没有提供三元组，直接使用LLM回答
        logger.info("没有提供三元组，直接使用LLM回答")
        response = LLM.complete(f"回答问题：{question}")
        result = response.text if hasattr(response, 'text') else str(response)
        logger.debug(f"LLM回答: {result}")
        return result
        
    # 构建知识图谱索引
    logger.info(f"使用 {len(triples)} 个三元组构建知识图谱索引")
    kg_index = _build_lindex(triples)
    
    # 使用知识图谱回答问题
    engine = kg_index.as_query_engine(llm=LLM, similarity_top_k=3)
    logger.info("使用知识图谱查询引擎回答问题")
    response = engine.query(question)
    result = response.response if hasattr(response, 'response') else str(response)
    logger.debug(f"知识图谱回答: {result}")
    
    return result
